// drvload.cpp : Defines the entry point for the console application.
//

#include "stdafx.h"
//  #include <ntstatus.h>
#include "utility.h"
#include "resource.h"
#include <iostream>

#define SE_LOAD_DRIVER_PRIVILEGE     (10L)  //  FROM  WDM.H
#define CURRENT_IDR     IDR_DRIVER          //  IDR_ALTEREG
#define USER_LIMIT  4

  //  space, hyphen, some slash's
const wchar_t strTrimChars[] = { 0x20, 0x2D, 0x2F, 0x5C, 0xC4, 0x0 };

inline bool isComplex(LPCWSTR astr)
{
    if(::StrChrIW(astr, '\\')) return true;
    if(::StrChrIW(astr, '/')) return true;
    return false;
}

bool setPathbin(LPCWSTR aval, BSTR * presult, BSTR * psuffix = NULL)
{
   * presult = NULL;
   CComBSTR buffer(MAX_PATH + 1);
      * buffer = 0;
   LPWSTR   suffix = NULL;
   DWORD retval = ::GetFullPathNameW(aval, MAX_PATH, buffer, & suffix);
   if(!(* buffer) || (!suffix))  retval = 0;

   if(retval)
     if(* suffix)
     {
       * presult = buffer.Detach();
       if(psuffix) * psuffix = ::SysAllocString(suffix);
     }

   return ((* presult) != NULL);
}

int wmain(int argc, wchar_t* argv[])
{
  if(!(argc >> 1))  //  parahelp
  {
  std::cout << "Application loads driver that will be registered only for current session, up to reboot.\n";
  std::cout << "Required parameter is driver file.\n";
  std::cout << "Optional parameter is registry alias.\n";
  std::cout << "At the end switch <<-u>> to avoid loading driver to kernel, if required only unloading.\n";
  std::cout << "You can use <<sc start/query/stop %alias%>> since registration in cruising regime.\n";
  std::cout << std::endl;
	    return 0;
  }

  bool  thehalf(false); //  only  unload driver
  DWORD retval(ERROR_FILE_NOT_FOUND);
  CComBSTR  thealias, theimage;
const _ATL_REGMAP_ENTRY strRgsvars = { L"APPID", L"MODULE" };

  try {
    LPCWSTR switches = &(strTrimChars[1]);
    _ATL_REGMAP_ENTRY candidates = { NULL, NULL };
    int  ambiguous(1);
    for(int i(1), n(USER_LIMIT); (i < argc) && (n); ++i, --n)
    {
        LPCWSTR thestr = argv[i];
        if(::PathFileExistsW(thestr))
        {
            bool  theheap(false), thestack(false);
            if(candidates.szData)
            {
                theheap = isComplex(candidates.szData);
                thestack= isComplex(thestr);
                if((!theheap && thestack) || (theheap && !thestack))
                    if(thestack) std::swap(candidates.szData, thestr);
            }

            if((theheap && thestack) || (!theheap && !thestack))
            {
                candidates.szData = thestr;
                --ambiguous;
                continue;
            };
        };

        if(::StrChrIW(switches, *thestr))
        {
           int  n(USER_LIMIT);
           for(++thestr; (*thestr) && (n); ++thestr, --n)
              if(!::StrChrIW(switches, *thestr))
                  break;
           if((*thestr == 'u') || (*thestr == 'U'))
                  thehalf = true;
        }
        else
          if(!candidates.szKey) candidates.szKey = thestr;
    };

    if((argc ^ 3) || (!thehalf) || (!candidates.szKey))
      if(ambiguous < 0)
      {
          std::cout << "Parameters are ambiguous.\n";
          throw DWORD(ERROR_AMBIGUOUS_SYSTEM_DEVICE);
      }
      else
      if(ambiguous)
      {
          std::cout << "Path not found from arguments.\n";
          throw DWORD(ERROR_FILE_NOT_FOUND);
      };

    if(candidates.szKey)
    {
        if(candidates.szData)
        {
           if(setPathbin(candidates.szData, & theimage))
                 retval = ERROR_SUCCESS;
        }
        else retval = ERROR_SUCCESS;
        thealias = candidates.szKey;
    }
    else
        if(setPathbin(candidates.szData, & theimage, & thealias))
            retval = ERROR_SUCCESS;

    if(candidates.szData)
        if(::PathIsDirectoryW(theimage)) 
            retval = ERROR_FILE_NOT_FOUND;
    if(!(thealias.Length() >> 1))    retval = ERROR_BAD_LENGTH;
       else if(thealias.Length() > MAX_PATH)
                retval = ERROR_LABEL_TOO_LONG;

    if(retval ^ ERROR_SUCCESS)
        std::cout << "Some parameters are not accepted.\n";
    else 
    if(candidates.szData)
    {
        Cdummodule::AddReplacements(strRgsvars.szKey, thealias.operator BSTR());
        Cdummodule::AddReplacements(strRgsvars.szData, theimage.operator BSTR());
    };
  }
  catch(DWORD err)
  {
      retval = err;
  }
  catch(...)
  {
      retval = E_UNEXPECTED;
  };

  if(retval ^ ERROR_SUCCESS)
  {
 std::cout << "The process is terminating abnormally, the stage of the parsing of the user input have ended up as a bit catastrophic.\n";
 std::cout << "\treturned error code:\t0x" << std::hex << retval << std::endl;
      return retval;
  }

  {
     std::cout << "\n\tThe summary of the job.\n";
     std::wcout << L"The driver name:\t" << thealias.operator BSTR() << std::endl;
     if(theimage.Length())
        std::wcout << L"The binary file:\t" << theimage.operator BSTR() << std::endl;
     else
        std::cout << "The registration operation will be skipped.\n";
  }

  retval = Cdummodule::IncludePrivilege(SE_LOAD_DRIVER_PRIVILEGE);
  if(retval ^ ERROR_SUCCESS)
  {
     std::cout << "\nPrivileges called <<Load device drivers>> was not activated in your environment.\n";
     std::cout << "Perhaps, you aren't memeber of any local administrator's, check setting's by 'secpol.msc'.\n";
     std::cout << "\treturned error code:\t0x" << std::hex << retval << std::endl;
     return retval;
  }

  CComBSTR  theregkey;
      theregkey.LoadStringW(IDS_REGSUBKEY);
      theregkey.AppendBSTR(thealias);

  Cdummodule atlModule(theregkey);
  bool theblank = atlModule.IsEmpty();

  if(!(theimage.Length()) && theblank)
  {
      std::cout << "\nThe target driver does not exist in the registry.\n";
      std::cout << "The process is terminating abnormally.\n";
      if(thehalf) retval = ERROR_INVALID_PARAMETER;
        else retval = ERROR_DATABASE_DOES_NOT_EXIST;
      std::cout << "\treturned error code:\t0x" << std::hex << retval << std::endl;
      return retval;
  }

  if(theimage.Length())
  try {
     if(theblank)
     {
       retval = atlModule.UncreateService(thealias);
       if((retval ^ ERROR_SUCCESS) && (retval ^ ERROR_SERVICE_EXISTS))
       {
           if(retval ^ ERROR_INVALID_NAME)
           {
 std::cout << "\nOne of SHDeleteKey, CreateService, OpenSCManager is fails,\n";
 std::cout << "\treturned error code:\t0x" << std::hex << retval << std::endl;
           }
           else {
 std::cout << "\nSome parameters are not accepted.\n";
 std::cout << "The filename, directory name, or volume label syntax is incorrect.\n";
           }
           throw retval;
       };

       if((!thehalf) || (retval ^ ERROR_SERVICE_EXISTS))
       {
           HRESULT hr = atlModule.UpdateRegistryFromResource(CURRENT_IDR);
           if(SUCCEEDED(hr))
           {
              std::cout << "\nTarget driver is recorded into registry:\n";
              std::wcout << theregkey.operator BSTR() << std::endl;
              std::cout << "This key will be automatically removed after reboot.\n";
           }
           else {
              std::cout << "\nRegistry is not updated, returned error code:\t0x" << std::hex ;
              std::wcout << hr << std::endl << theregkey.operator BSTR() << std::endl;
           };
       }
       else {
          std::cout << "\nTarget driver already exists:\n\t";
          std::wcout << thealias.operator BSTR() << std::endl;
          std::cout << "Nothing change is made in registry.\n";
       };
     }
     else {
        std::cout << "\nTarget registry key already exists:\n";
        std::wcout << theregkey.operator BSTR() << std::endl;
        std::cout << "Nothing change is made there.\n";
        if(!thehalf)
        if(atlModule.SearchAmbiguous(CURRENT_IDR))
        {
            std::cout << "\nIt looks like that subject registry key was created not this application,\n";
            std::cout << "you can control them by system tool 'sc.exe', including to delete them.\n";
            retval = ERROR_AMBIGUOUS_SYSTEM_DEVICE;
        };
     };
  }
  catch(DWORD err)
  {
      retval = err;
  }
  catch(...)
  {
      retval = E_UNEXPECTED;
  };

  if(retval ^ ERROR_SUCCESS)
  {
     std::cout << "The process is terminating abnormally, the registration stage is a bit catastrophic.\n";
     return retval;
  }

  retval = atlModule.UpdateLoadDriver(thehalf);
  if(retval ^ ERROR_SUCCESS)
  {
      std::cout << "\nTarget driver was not loaded,\n";
      std::cout << "operation returned error code:\t0x" << std::hex << retval << std::endl;
  }
  else
     if(thehalf)
     {
        if(!theblank) std::cout << "Target driver is successfully unloaded.\n";
     }
     else std::cout << "Target driver is successfully loaded.\n";

	return retval;
}

